from ARM_TGNCF.Parse_ARM_TGNCF import ARM_args, inter_scores
from loaddata.load_data_tool.load_events_data import load_events_data


def load_ratings_data():
    data_all = load_events_data()[['userId', 'itemId', 'rating']]

    ratings_num = len(data_all)

    if ARM_args.forecast_method == 'inter':

        inter_score = inter_scores[ARM_args.dataset]
        ratings = list(data_all['rating'])
        for index in range(ratings_num):
            if ratings[index] < inter_score:
                data_all.loc[index, 'rating'] = 0
            else:
                data_all.loc[index, 'rating'] = 1

    train_len = int(ratings_num * ARM_args.train)
    train_data = data_all[:train_len]
    test_data = data_all[train_len:]

    return train_data, test_data, ratings_num
